# -*- coding: utf-8 -*-
"""
Created on Wed Jul 15 16:22:36 2020

@author: Administrator
"""

import avadata as avadata
import cv2
import numpy as np
from collections import defaultdict
from arg import cmd
from tqdm import tqdm
from fvcore.common.file_io import PathManager
import os

class videoReader(object):
    def __init__(self, source):
        self.source=source
        self.cap=cv2.VideoCapture(self.source)
        self.display_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.display_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.fps=round(self.cap.get(cv2.CAP_PROP_FPS))
        if all_path.ava_on ==1:
            
            self.cap.set(cv2.CAP_PROP_POS_FRAMES,902*self.fps) #从902秒开始读
        self.valid_frames=int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.endtime=self.valid_frames
        if all_path.endsec:
            self.endtime=endsec*self.fps
        
        self.output_file=None
        if all_path.outpath:
            self.output_file=self.get_output_file(all_path.outpath)#r'C:/Users/Administrator/Downloads/haikang/ava/output.mp4'
        
        if not self.cap.isOpened():
            raise IOError("Video {} cannot be opened".format(self.source))
    def __iter__(self):
        return self
    
    def __next__(self):
        pos = self.cap.get(cv2.CAP_PROP_POS_FRAMES)#获取当前帧号
        was_read,frame =self.cap.read()
        
        if pos < self.endtime:
            
            return was_read, frame, pos
        else:
            raise StopIteration
    
        if not was_read :
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
            frame = None
        return was_read, frame, pos
    def get_output_file(self, path):
        return cv2.VideoWriter(
            filename=path,
            fourcc=cv2.VideoWriter_fourcc(*"mp4v"),
            fps=float(25),
            frameSize=(self.display_width, self.display_height),
            isColor=True,
        )
        
    def display(self, frame):

        if self.output_file is None:
            cv2.imshow("SlowFast", frame)
        else:
            self.output_file.write(frame)
       
    def clean(self):
        self.cap.release()
        
        if self.output_file is None:
            cv2.destroyAllWindows()
        else:
            self.output_file.release()
        
def getitem_by_name(name, ava_on ,all_box,collect,mode_on):# 运行文件后执行这个函数，注意要选一个训练集里有的视频
    '''
    arg:
        name of the video (e.g.'BCiuXAuCKAU')
        only label the keyframe if mode_on=True else label every frames according to the sec they belong to. 
    return:
        no return but will output a frame with box and label 
    '''
    if ava_on == 1:
        all_box=avadata.load_boxes_and_labels('train' )
        name_ele=name.split('/')[-1].split('.')[0] #从路径中获取视频的名称
        
        assert name_ele in all_box.keys( )
        labellist=[]
        
        with open(path_label) as f:
            for i, ele in enumerate( f.readlines()):
                if i % 5 ==1:
                    labellist.append(ele.split('"')[1])
                                    
        item=all_box[name_ele]
        
        frame_provider=videoReader(name)
    else:
        frame_provider = videoReader(os.path.join(path_video,name))
        
        assert name in all_box.keys()
        item=all_box[name]
    
    
    box_sec=defaultdict(list)
    label_sec=defaultdict(list)
    
    
    for sec in item.keys():
        if len(item[sec]) ==0:
            box_sec[sec]
            label_sec[sec]
        else:
            for ele in item[sec]:
                if ava_on == 1:
                    box_sec[sec].append(ele[0])
                    label_sec[sec].append(list(map(int,ele[1])))
                else:
                    box_sec[sec].append(ele[0])
                    label_sec[sec].append(ele[1])
    
    frameset=[]
    for able_to_read, frame, pos in tqdm(frame_provider,desc='processing frame'): 
        
        if not able_to_read:
            frame=[]
            break
        
        index=pos if mode_on else int(pos // frame_provider.fps)
        if len(item[index])==0: #or之前是为了只框选关键帧，or之后是为了处理没有被label的关键帧
            if collect==1:
                frameset.append(frame)
                continue
            else:
                frame_provider.display(frame)
                continue
        else:
            
                
            boxes=box_sec[index]
            labelset=label_sec[index]
            
    
            boxes=np.asarray(boxes)
            if ava_on == 1:
                boxes[:,[0,2]]*=frame_provider.display_width
                boxes[:,[1,3]]*=frame_provider.display_height
            boxes=clip_boxes_to_image(boxes, frame_provider.display_height, frame_provider.display_width)
            labelset=np.asarray(labelset)
            
            
            
            
            for i,(box, box_labels) in enumerate(zip(boxes.astype(int), labelset)):
                cv2.rectangle(
                    frame,
                    tuple(box[:2]),
                    tuple(box[2:]),
                    (0,255,0),
                    thickness=2,
                )
                
                label_origin = box[:2]
                label_origin[-1]-= 10*i
                for label in box_labels:
                    if ava_on:
                        labelname=labellist[label-1]
                    else:
                        labelname=label
                    (label_width, label_height), _ = cv2.getTextSize(
                        labelname, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
                    )
                    label_origin[-1] -=  5
                    
                    cv2.putText(
                        frame,
                        labelname,
                        tuple(label_origin),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.8,
                        (0, 0, 255),
                        2,
                    )
                    label_origin[-1] -= label_height + 5
                    
        if collect == 1 :
            
            if mode_on:           
                for i in range(30):
                    frameset.append(frame)
            else:
                frameset.append(frame)
        else:
            if mode_on:           
                 for i in range(30):
                     frame_provider.display(frame)
            else:
                 frame_provider.display(frame)
        #key = cv2.waitKey(1)
        #if key == 27:
            #break
    frame_provider.clean()
    if collect ==1 :
        return frameset

def collectvideo(all_box,path,mode_on,path_video):
    
    all_frame=[]
    for name in all_box.keys():
        all_frame.extend(getitem_by_name(name, ava_on,all_box ,collect ,mode_on))
    writer=cv2.VideoWriter(
            filename=path,
            fourcc=cv2.VideoWriter_fourcc(*"mp4v"),
            fps=float(30),
            frameSize=(1920, 1080),
            isColor=True,
        )
    for frame in all_frame:
        writer.write(frame)

def load_box_by_name(path_annotation, valid_frames=range(30)):
    
    filename=path_annotation
    all_boxes={}
    
    with PathManager.open(filename,'r') as f:
        for i, line in enumerate(f):
            if i ==0:
                continue
            else :
                row = line.strip().split(',')
                nt=row[0].split('"')[1].split('#t=')
                video_name, frame_sec = nt[0], int(float(nt[1]))#row[0], int(row[1])
                    
    
                # Box with format [x1, y1, x2, y2] with a range of [0, 1] as float.
                box_key = ",".join(row[1:5])
                box = list(map(float, row[1:5]))
                
    
                if video_name not in all_boxes:
                    all_boxes[video_name] = {}
                    for sec in valid_frames:
                        all_boxes[video_name][sec] = {}
    
                if box_key not in all_boxes[video_name][frame_sec]:
                    all_boxes[video_name][frame_sec][box_key] = [box, []]
                    
    
                all_boxes[video_name][frame_sec][box_key][1].append(row[5].split('"')[1])
                
    for video_name in all_boxes.keys():
        for frame_sec in all_boxes[video_name].keys():
            # Save in format of a list of [box_i, box_i_labels].
            all_boxes[video_name][frame_sec] = list(
                all_boxes[video_name][frame_sec].values()
            )
    return all_boxes

def load_box_by_name2(path_annotation, valid_frames=range(900)):
    
    filename=path_annotation
    all_boxes={}
    
    with PathManager.open(filename,'r') as f:
        for i, line in enumerate(f):
            if i ==0:
                continue
            else :
                row = line.strip().split(',')
                nt=row[0].split('"')[1].split('#t=')
                video_name, frame_sec = nt[0], round(float(nt[1])*25)#row[0], int(row[1])
                    
    
                # Box with format [x1, y1, x2, y2] with a range of [0, 1] as float.
                box_key = ",".join(row[1:5])
                box = list(map(float, row[1:5]))
                
    
                if video_name not in all_boxes:
                    all_boxes[video_name] = {}
                    for sec in valid_frames:
                        all_boxes[video_name][sec] = {}
    
                if box_key not in all_boxes[video_name][frame_sec]:
                    all_boxes[video_name][frame_sec][box_key] = [box, []]
                    
    
                all_boxes[video_name][frame_sec][box_key][1].append(row[5].split('"')[1])
                
    for video_name in all_boxes.keys():
        for frame_sec in all_boxes[video_name].keys():
            # Save in format of a list of [box_i, box_i_labels].
            all_boxes[video_name][frame_sec] = list(
                all_boxes[video_name][frame_sec].values()
            )
    return all_boxes


def clip_boxes_to_image(boxes, height, width):

    """

    Clip the boxes with the height and width of the image size.

    Args:

        boxes (ndarray): bounding boxes to peform crop. The dimension is

        `num boxes` x 4.

        height (int): the height of the image.

        width (int): the width of the image.

    Returns:

        boxes (ndarray): cropped bounding boxes.

    """

    boxes[:, [0, 2]] = np.minimum(

        width - 1.0, np.maximum(0.0, boxes[:, [0, 2]])

    )

    boxes[:, [1, 3]] = np.minimum(

        height - 1.0, np.maximum(0.0, boxes[:, [1, 3]])

    )

    return boxes

if  __name__=="__main__":
    
    all_path=cmd()

    path_label=all_path.label 
    
    videoname= all_path.name    
    endsec=all_path.endsec
    path_video=all_path.video
    path_annotation=all_path.annotation
    ava_on=all_path.ava_on
    collect=all_path.collect
    mode_on=True
    
    if mode_on:
        all_box=all_box=load_box_by_name2(path_annotation)
    else:
        all_box=all_box=load_box_by_name(path_annotation)
    if collect ==1:
        
        collectvideo(all_box,all_path.outpath,mode_on,path_video)
    
    else:        
        
        
        getitem_by_name(videoname,ava_on, all_box, collect, mode_on)
           
